%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Explanation of what the file does
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% This function solves for a mobility matrix across periods and a labor 
% supply matrix across periods given a guess for the path of utility hats
% as well as an initial mobility matrix and an initial labor supply
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Inputs needed
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% 1. Guess of path of utility hats (Uguess), size I*S*T
% 2. Mobility matrix in baseline economy (MuBas), size (I*S)*(I*S)*T
% 3. Initial labor supply matrix (EllInit), size I*S
% 4. Discount factor (ba), scalar
% 5. Migration elasticity (ka), scalar
% 6. Sectoral elasticity (nu), scalar
% 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Outputs produced
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% 1. Mobility matrix in counterfactual economy (MuEra), size (I*S)*(I*S)*T
% 2. Labor supply matrix in ocunterfactual (EllEra), size I*S*T
% 3. The theta zero matrix for utility updating (TZ), size (I*S)*(I*S)

function [MuEra,EllEra,TZ] = NECO_Block1(Uguess,MuBas,EllInit,ba,ka,nu)

% Define some important variables

I                 = size(EllInit,1);
S                 = size(EllInit,2);
T                 = size(Uguess,3);
MuEra             = NaN(I*S,I*S,T);
EllEra            = NaN(I,S,T);
MuHashB           = NaN(I*S,I,T);
MuCondB           = NaN(size(MuBas));
for t=1:T
    MuHashB(:,:,t)= MuBas(:,:,t) * kron(eye(I),ones(S,1));
    MuCondB(:,:,t)= MuBas(:,:,t) ./ kron(MuHashB(:,:,t),ones(1,S));
end
MuCondB(isnan(MuCondB)) = 0;
MuHashDB          = MuHashB(:,:,2:end)./MuHashB(:,:,1:end-1);
MuHashDB(isnan(MuHashDB)) = 0;
MuCondDB          = MuCondB(:,:,2:end)./MuCondB(:,:,1:end-1);
MuCondDB(isnan(MuCondDB)) = 0;

% Period zero mobility and labor

MuEra(:,:,1)      = MuBas(:,:,1);
EllEra(:,:,1)     = EllInit;

% Period one labor

EllInter          = EllEra(:,:,1)';
EllNewV           = MuEra(:,:,1)' * EllInter(:);
EllEra(:,:,2)     = reshape(EllNewV,S,I)';

% Period one mobility

UdotInter         = Uguess(:,:,1)';
TZ                = (repmat(UdotInter(:)',I*S,1)).^(ba/nu)...
                    .* MuCondB(:,:,2);
UdotInter         = Uguess(:,:,2)';
MuCondNum         = (repmat(UdotInter(:)',I*S,1)).^(ba/nu) .* TZ;
MuCondDen         = kron(MuCondNum*kron(eye(I),ones(S,1)),ones(1,S));
MuCond            = MuCondNum./MuCondDen;
MuCond(isnan(MuCond)) = 0;
MuHashNum         = MuHashB(:,:,2) .* (MuCondNum*kron(eye(I),ones(S,1)))...
                    .^(nu/ka);
MuHashDen         = kron(sum(MuHashNum,2),ones(1,I));
MuHash            = MuHashNum./MuHashDen;
MuEra(:,:,2)      = MuCond .* kron(MuHash,ones(1,S));

% Run the forward loop where the mus of a period are found using the mus of
% the previous period and the u hats of the following period, as well as 
% the mu dots of the baseline economy. Obtain labor supplies from previous 
% period labor supplies and the mus

for t=3:T
    UdotInter     = Uguess(:,:,t)';
    MuHashPre     = MuEra(:,:,t-1) * kron(eye(I),ones(S,1));
    MuCondPre     = MuEra(:,:,t-1) ./ kron(MuHashPre,ones(1,S));
    MuCondPre(isnan(MuCondPre)) = 0;
    MuCondNum     = (repmat(UdotInter(:)',I*S,1)).^(ba/nu) .* MuCondPre...
                    .* MuCondDB(:,:,t-1);
    MuCondDen     = kron(MuCondNum*kron(eye(I),ones(S,1)),ones(1,S));
    MuCond        = MuCondNum./MuCondDen;
    MuCond(isnan(MuCond)) = 0;
    MuHashNum     = MuHashPre .* MuHashDB(:,:,t-1) .* (MuCondNum*...
                    kron(eye(I),ones(S,1))).^(nu/ka);
    MuHashDen     = kron(sum(MuHashNum,2),ones(1,I));
    MuHash        = MuHashNum./MuHashDen;
    MuEra(:,:,t)  = MuCond .* kron(MuHash,ones(1,S));
    EllInter      = EllEra(:,:,t-1)';
    EllNewV       = MuEra(:,:,t-1)' * EllInter(:);
    EllEra(:,:,t) = reshape(EllNewV,S,I)';
end

end
